package com.leetcode.tree;

import com.leetcode.basic.TreeNode;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author Dennis Li
 * @date 2020/12/25 16:56
 */
public class DelNodes_1110 {

    public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        Set<Integer> set = Arrays.stream(to_delete)
                .boxed().collect(Collectors.toSet());
        List<TreeNode> list = new ArrayList<>();

        addToList(list, root, set);
        return list;
    }

    public void delNodes(List<TreeNode> list, TreeNode root, Set<Integer> toDelete) {
        if (root == null) {
            return;
        }
        // 保证root是不在toDelete中的
        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        while (!queue.isEmpty()) {
            int size = queue.size();
            while (size-- > 0) {
                TreeNode node = queue.poll();
                if (helper(list, toDelete, queue, node.left)) {
                    node.left = null;
                }
                if (helper(list, toDelete, queue, node.right)) {
                    node.right = null;
                }
            }
        }
    }

    private boolean helper(List<TreeNode> list, Set<Integer> toDelete, Queue<TreeNode> queue, TreeNode node) {
        if (node != null) {
            if (toDelete.contains(node.val)) {
                addToList(list, node.left, toDelete);
                addToList(list, node.right, toDelete);
                return true;
            } else {
                queue.offer(node);
            }
        }
        return false;
    }

    private void addToList(List<TreeNode> list, TreeNode node, Set<Integer> toDelete) {
        if (node != null) {
            if (toDelete.contains(node.val)) {
                addToList(list, node.left, toDelete);
                addToList(list, node.right, toDelete);
            } else {
                list.add(node);
                delNodes(list, node, toDelete);
            }
        }
    }

}
